from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_ollama import ChatOllama
import json
from langchain_core.messages import ToolMessage
from langchain_core.tools import tool
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.types import Command, interrupt
from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()

class State(TypedDict):
    # Messages have the type "list". The `add_messages` function
    # in the annotation defines how this state key should be updated
    # (in this case, it appends messages to the list, rather than overwriting them)
    messages: Annotated[list, add_messages]

llm = ChatOllama(
    model="qwen2.5:7b-instruct-q5_K_S",
    temperature=0.2,
)

@tool
def get_weather(city: str) -> str:
    """Get weather for a given city."""
    print(f"Getting weather for {city}")
    return f"It's always raining in {city}!"

@tool
def human_assistance(query: str) -> str:
    """Request assistance from a human."""
    print(f"Requesting human assistance for: {query}")
    human_response = interrupt({"query": query})
    return human_response["data"]

tools = [get_weather, human_assistance]

llm_with_tools = llm.bind_tools(tools)

# chatbot node
def chatbot(state: State):
    message = llm_with_tools.invoke(state["messages"])
    # Because we will be interrupting during tool execution,
    # we disable parallel tool calling to avoid repeating any
    # tool invocations when we resume.
    assert len(message.tool_calls) <= 1
    return {"messages": [message]}

tool_node = ToolNode(tools=tools)

graph_builder = StateGraph(State)

# The first argument is the unique node name
# The second argument is the function or object that will be called whenever
# the node is used.
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node("tools", tool_node)
# The `tools_condition` function returns "tools" if the chatbot asks to use a tool, and "END" if
# it is fine directly responding. This conditional routing defines the main agent loop.
graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
)
# Any time a tool is called, we return to the chatbot to decide the next step
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_edge(START, "chatbot")

graph = graph_builder.compile(checkpointer=memory)

my_config = {"configurable": {"thread_id": "1"}}

human_response = (
    "We, the experts are here to help! We'd recommend you check out LangGraph to build your agent."
    " It's much more reliable and extensible than simple autonomous agents."
)

human_command = Command(resume={"data": human_response})

def stream_graph_updates(user_input: str):
    my_input = None
    if user_input.lower() == "human":
        my_input = human_command
    else:
        my_input = {"messages": [{"role": "user", "content": user_input}]}

    for event in graph.stream(input=my_input,
                              config=my_config,
                              stream_mode="values",
                              ):
        for value in event.values():
            # print(f"Assistant: {value[-1].content}")
            event["messages"][-1].pretty_print()

while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break
    
        stream_graph_updates(user_input)
    except:
        # fallback if input() is not available
        user_input = "What do you know about LangGraph?"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break

